import random

import cv2
import numpy as np
import time


class CourseGenerator():

    def __int__(self):
        self.sk_list = None
        self.time_ = ''
        self.id_dict = {}
        self.skeleton_path = ''

    def get_keyFrame(self, file_path):
        self.frame_info = {}
        if file_path == "" or file_path.split('.')[-1] not in ['wmv', 'avi', 'mp4', 'mov']:
            return
        self.file_path = file_path
        file_name = file_path.split('/')[-1]
        self.cap = cv2.VideoCapture(file_path)
        _, self.icon = self.cap.read()
        self.sk_list = None
        timestamp = str(time.time())
        self.icon_path = '/static/upload/action/' + timestamp + '.png'
        cv2.imwrite('.' + self.icon_path, self.icon)
        (frame, skeleton), self.sk_list, self.sk_list1 = get_keyFrame(self.file_path)
        for i in range(len(frame)):
            # 转毫秒
            frame_ms = int(frame[i] / self.cap.get(cv2.CAP_PROP_FPS) * 1000)
            # 默认没有描述，默认关注所有关节
            self.frame_info[frame_ms] = ['', 1023, skeleton[i]]
        skeleton_list = np.stack(self.sk_list1)
        self.skeleton_path = '/static/upload/course_action/skeleton/' + file_name[:file_name.rfind('.')] + '.npy'
        np.save('.' + self.skeleton_path, skeleton_list)


def keyframe_extract(skeleton_list, visibility=0.8, movements=0.09):
    skeleton_list = np.asarray(skeleton_list)
    frames = skeleton_list.shape[0]
    kps = skeleton_list.shape[1]
    visib = skeleton_list.shape[2]

    key_frames = []
    skeleton_filter_list = []

    for frame in range(frames - 2):
        skeleton_i = skeleton_list[frame]
        skeleton_ii = skeleton_list[frame + 1]

        diff = np.ones([kps, 2])
        diff[:, 0] = np.linalg.norm((skeleton_i[:, 0:3] - skeleton_ii[:, 0:3]), axis=1)
        if visib == 4:
            diff[:, 1] = skeleton_i[:, 3] * skeleton_ii[:, 3]

        # move_flag = 0

        for idx in range(kps):
            if diff[idx, 1] > visibility:
                if diff[idx, 0] > movements:
                    # move_flag = 1
                    key_frames.append(frame)
                    skeleton_filter_list.append(skeleton_i)
                    break
        if frame == 0:
            key_frames.append(frame)
            skeleton_filter_list.append(skeleton_i)

    return key_frames, skeleton_filter_list


def get_keyFrame(video_path):
    skeleton_list = []
    skeleton_list1 = []
    cap = cv2.VideoCapture(video_path)
    success, img = cap.read()
    skeleton = np.zeros((33, 3))
    skeleton1 = np.zeros((33,4))
    import mediapipe as mp
    pose = mp.solutions.pose.Pose(static_image_mode=False,
                                  model_complexity=2,
                                  smooth_landmarks=True,
                                  enable_segmentation=True,
                                  min_detection_confidence=0.5,
                                  min_tracking_confidence=0.5)
    while success:
        img_RGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        results = pose.process(img_RGB)
        if results.pose_world_landmarks:
            for i in range(33):
                skeleton[i][0] = results.pose_world_landmarks.landmark[i].x
                skeleton[i][1] = results.pose_world_landmarks.landmark[i].y
                skeleton[i][2] = results.pose_world_landmarks.landmark[i].z
            for i in range(33):
                skeleton1[i][0] = results.pose_world_landmarks.landmark[i].x
                skeleton1[i][1] = results.pose_world_landmarks.landmark[i].y
                skeleton1[i][2] = results.pose_world_landmarks.landmark[i].z
                skeleton1[i][3] = results.pose_world_landmarks.landmark[i].visibility
        skeleton_list.append(skeleton_tran(skeleton))
        skeleton_list1.append(skeleton_tran(skeleton1))
        success, img = cap.read()

    return keyframe_extract(skeleton_list), skeleton_list, skeleton_list1


def skeleton_tran(skeleton, norm=False):
    # blazepose 33-> 15

    n = [-1, 24, 26, 28,
         23, 25, 27,
         12, 14, 16,
         11, 13, 15,
         -1, 0]
    skeleton_ = np.zeros((15, skeleton.shape[-1]))
    for i in range(3):
        skeleton_[0][i] = (skeleton[23][i] + skeleton[24][i]) / 2
        skeleton_[13][i] = (skeleton[11][i] + skeleton[12][i]) / 2
    for i in range(15):
        if n[i] != -1:
            skeleton_[i] = skeleton[n[i]]
    if norm:
        a = skeleton_[0].copy()
        for i in range(15):
            skeleton_[i] -= a
    return skeleton_
